import numpy as np


class NTUDGraph:
    def __init__(self, labeling_mode='spatial'):
        num_node = 25
        self_link = [(i, i) for i in range(num_node)]
        inward_ori_index = [(1, 2), (2, 21), (3, 21), (4, 3), (5, 21), (6, 5),
                            (7, 6), (8, 7), (9, 21), (10, 9), (11, 10),
                            (12, 11), (13, 1), (14, 13), (15, 14), (16, 15),
                            (17, 1), (18, 17), (19, 18), (20, 19), (22, 23),
                            (23, 8), (24, 25), (25, 12)]
        inward = [(i - 1, j - 1) for (i, j) in inward_ori_index]
        outward = [(j, i) for (i, j) in inward]
        neighbor = inward + outward

        self.num_node = num_node
        self.self_link = self_link
        self.inward = inward
        self.outward = outward
        self.neighbor = neighbor
        self.A = self.get_adjacency_matrix(labeling_mode)

    def get_adjacency_matrix(self, labeling_mode=None):
        if labeling_mode is None:
            return self.A
        if labeling_mode == 'spatial':
            A = get_spatial_graph(self.num_node, self.self_link, self.inward, self.outward)
        else:
            raise ValueError()
        return A


class MulScale_NTUDGraph:
    def __init__(self, max_hop=2, dilation=1):
        self.max_hop = max_hop
        self.dilation = dilation
        self.get_mul_graph()

        self.A_1 = get_hops_graph(self.graph_1, self.num_node, max_hop, dilation)
        self.A_2 = get_hops_graph(self.graph_2, len(self.parts_2), max_hop, dilation)
        self.A_3 = get_hops_graph(self.graph_3, len(self.parts_3), max_hop, dilation)

        self.mask_norm_2 = normalize_digraph(self.mask_2)
        self.mask_norm_3 = normalize_digraph(self.mask_3)

    def get_mul_graph(self):
        num_node = 25
        self_link = [(i, i) for i in range(num_node)]
        inward_ori_index = [(1, 2), (2, 21), (3, 21), (4, 3), (5, 21), (6, 5),
                            (7, 6), (8, 7), (9, 21), (10, 9), (11, 10),
                            (12, 11), (13, 1), (14, 13), (15, 14), (16, 15),
                            (17, 1), (18, 17), (19, 18), (20, 19), (22, 23),
                            (23, 8), (24, 25), (25, 12)]
        inward = [(i - 1, j - 1) for (i, j) in inward_ori_index]
        outward = [(j, i) for (i, j) in inward]
        neighbor = inward + outward

        self.num_node = num_node
        self.self_link = self_link
        self.inward = inward
        self.outward = outward
        self.neighbor = neighbor
        self.edge = self_link + neighbor

        self.graph_1 = edge2mat(self.edge, num_node)

        self.parts_2 = [
            np.array([1]) - 1,  # torso
            np.array([2, 21]) - 1,
            np.array([3, 4]) - 1,  # head
            np.array([5, 6, 7]) - 1,  # left_arm
            np.array([8, 22, 23]) - 1,  # left_hand
            np.array([9, 10, 11]) - 1,  # right_arm
            np.array([12, 24, 25]) - 1,  # right_hand
            np.array([13, 14]) - 1,  # left_leg
            np.array([15, 16]) - 1,  # left_foot
            np.array([17, 18]) - 1,  # right_leg
            np.array([19, 20]) - 1,  # right_foot
        ]
        self.mask_2 = self.get_marks(self.num_node, self.parts_2) # [25,11]
        # graph_2 = M^t * graph_1 * M = [11,25]*[25*25]*[25,11] = [11,11]
        graph_2 = np.dot(np.dot(self.mask_2.T, self.graph_1), self.mask_2)
        self.graph_2 = np.where(graph_2, np.ones_like(graph_2), graph_2)

        self.parts_3 = [
            np.array([1, 2]) - 1,  # torso
            np.array([3]) - 1,  # head
            np.array([4, 5]) - 1,  # left_arm
            np.array([6, 7]) - 1,  # right_arm
            np.array([8, 9]) - 1,  # left_leg
            np.array([10, 11]) - 1,  # right_leg
        ]
        self.mask_3 = self.get_marks(len(self.parts_2), self.parts_3)
        # Same as graph_2, graph_3 = M^t * graph_2 * M = [6,11]* [11,11] * [11,6] =[6,6]
        graph_3 = np.dot(np.dot(self.mask_3.T, graph_2), self.mask_3)
        self.graph_3 = np.where(graph_3 > 1, np.ones_like(graph_3), graph_3)

    def get_marks(self, num_nodes, parts):
        num_parts = len(parts)
        part_array = np.zeros(shape=[num_nodes, num_parts],dtype='float')
        for i, part in enumerate(parts):
            for j in part:
                part_array[j, i] = 1
        return part_array


def edge2mat(link, num_node):
    A = np.zeros((num_node, num_node))
    for i, j in link:
        A[j, i] = 1
    return A


def normalize_digraph(A):
    Dl = np.sum(A, 0)
    h, w = A.shape
    Dn = np.zeros((w, w))
    for i in range(w):
        if Dl[i] > 0:
            Dn[i, i] = Dl[i] ** (-1)
    AD = np.dot(A, Dn)
    return AD


def get_spatial_graph(num_node, self_link, inward, outward):
    I = edge2mat(self_link, num_node)
    In = normalize_digraph(edge2mat(inward, num_node))
    Out = normalize_digraph(edge2mat(outward, num_node))
    A = np.stack((I, In, Out))
    return A


def get_hop_distance(num_node, A, max_hops):
    hop_dis = np.zeros((num_node, num_node)) + np.inf
    transfer_mat = [np.linalg.matrix_power(A, d) for d in range(max_hops + 1)]
    arrive_mat = (np.stack(transfer_mat) > 0)
    for d in range(max_hops, -1, -1):
        hop_dis[arrive_mat[d]] = d
    return hop_dis


def get_hops_graph(adj_matrix, num_node, max_hops=2, dilation=1):
    hop_dis = get_hop_distance(num_node, adj_matrix, max_hops)
    valid_hop = range(0, max_hops + 1, dilation)
    adjacency = np.zeros((num_node, num_node))
    for hop in valid_hop:
        adjacency[hop_dis == hop] = 1
    normalize_adjacency = normalize_digraph(adjacency)

    A = np.zeros((len(valid_hop), num_node, num_node))
    for i, hop in enumerate(valid_hop):
        A[i][hop_dis == hop] = normalize_adjacency[hop_dis == hop]
    return A
